import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
import proguard.gradle.ProGuardTask

import java.util.concurrent.atomic.AtomicInteger

buildscript {
    repositories {
        gradlePluginPortal()
        google()
    }
    dependencies {
        classpath "gradle.plugin.com.github.johnrengelman:shadow:${managedVersions['gradle.plugin.com.github.johnrengelman:shadow']}"
        classpath "com.guardsquare:proguard-gradle:${managedVersions['com.guardsquare:proguard-gradle']}"
    }
}

def relocatedProjects = projectsWithFlags('java', 'relocate')
def numConfiguredRelocatedProjects = new AtomicInteger()
configure(relocatedProjects) {
    // Generate the shaded JARs.
    task shadedJar(
            type: ShadowJar,
            group: 'Build',
            description: 'Builds the shaded main JAR.',
            dependsOn: tasks.classes) {

        configureShadowTask(project, delegate, true)
        archiveBaseName.set("${project.archivesBaseName}-shaded")

        // Exclude the class signature files.
        exclude '/META-INF/*.SF'
        exclude '/META-INF/*.DSA'
        exclude '/META-INF/*.RSA'
        // Exclude the files generated by Maven
        exclude '/META-INF/maven/**'
    }
    tasks.assemble.dependsOn tasks.shadedJar

    artifacts {
        archives shadedJar
    }

    task shadedClasses(
            type: Copy,
            group: 'Build',
            description: 'Extracts the shaded main JAR.',
            dependsOn: tasks.shadedJar) {

        from(zipTree(tasks.shadedJar.archivePath))
        from(sourceSets.main.output.classesDirs) {
            // Add the JAR resources excluded in the 'shadedJar' task.
            include '**/*.jar'
        }
        into "${project.buildDir}/classes/java/shaded-main"
    }

    task shadedTestJar(
            type: ShadowJar,
            group: 'Build',
            description: 'Builds the shaded test JAR.',
            dependsOn: tasks.testClasses) {

        configureShadowTask(project, delegate, false)
        archiveBaseName.set("test-${tasks.jar.archiveBaseName.get()}-shaded")
    }

    task copyShadedTestClasses(
            type: Copy,
            group: 'Build',
            description: 'Extracts the shaded test JAR.',
            dependsOn: tasks.shadedTestJar) {

        from(zipTree(tasks.shadedTestJar.archivePath))
        from(sourceSets.test.output.classesDirs) {
            // Add the JAR resources excluded in the 'shadedTestJar' task.
            include '**/*.jar'
        }
        into "${project.buildDir}/classes/java/shaded-test"
    }
}

// NB: Configure in a new closure so that all relocated projects have a 'shadedJar' task.
configure(relocatedProjects) {
    if (project.hasFlags('trim')) {
        // Task 'shadedJar' may produce a very large JAR. Rename it to '*-untrimmed-*.jar' and
        // let the task 'trimShadedJar' produce the trimmed JAR from it.
        tasks.shadedJar.archiveBaseName.set("${tasks.jar.archiveBaseName.get()}-untrimmed")

        task trimShadedJar(
                type: ProGuardTask,
                group: 'Build',
                description: 'Shrinks the shaded JAR by removing unused classes.') {

            relocatedProjects.each {
                dependsOn it.tasks.shadedJar.path
                dependsOn it.tasks.shadedTestJar.path
            }

            def shadedFile = tasks.shadedJar.archivePath
            def shadedAndTrimmedFile = file(shadedFile.path.replaceFirst('-untrimmed-', '-shaded-'))

            injars shadedFile
            // NB: By specifying 'outjars' *before* other 'injars' below, ProGuard will put only the classes
            //     from 'shadedFile' into 'shadedAndTrimmedFile'. See 'restructuring the output archives'
            //     for more information: https://www.guardsquare.com/en/proguard/manual/examples#restructuring
            outjars shadedAndTrimmedFile

            // Include all other shaded JARs so that ProGuard does not trim the classes and methods
            // that are used actually.
            injars tasks.shadedTestJar.archivePath
            relocatedProjects.each {
                if (it != project && !it.hasFlags('no_aggregation')) {
                    injars it.tasks.shadedJar.archivePath
                    injars it.tasks.shadedTestJar.archivePath
                }
            }

            def dependencyJars = new LinkedHashSet<File>()
            relocatedProjects.each { p ->
                // NB: ProGuardTask picks the dependencies added *after* evaluation correctly
                //     because libraryjar() intentionally keeps the specified dependencies as-is.
                //     See ProGuardTask.libraryjar() for more information.
                if (!p.hasFlags('no_aggregation')) {
                    doFirst {
                        p.configurations.testRuntimeClasspath.collect().each { File file ->
                            if (!file.path.startsWith("${rootProject.projectDir}")) {
                                dependencyJars.add(file)
                            }
                        }
                    }
                }
            }
            libraryjars files(dependencyJars)

            // Add JDK. Use modules instead of rt.jar if possible.
            File jmodDir = file("${System.getProperty('java.home')}/jmods")
            if (jmodDir.isDirectory()) {
                jmodDir.listFiles().findAll { File f ->
                    f.isFile() && f.name.toLowerCase(Locale.ENGLISH).endsWith(".jmod")
                }.sort().each { libraryjars it }
            } else {
                libraryjars file("${System.getProperty('java.home')}/lib/rt.jar")
            }

            dontoptimize
            dontobfuscate
            dontwarn // Ignore the harmless 'missing classes' warnings related with the optional dependencies.

            keepattributes '**'
            keepparameternames

            printconfiguration file("${project.buildDir}/proguard.cfg")
        }

        tasks.shadedClasses.dependsOn tasks.trimShadedJar
        tasks.assemble.dependsOn tasks.trimShadedJar

        // Add the trimmed JAR to archives.
        artifacts {
            trimShadedJar.outJarFiles.each {
                archives it
            }
        }
    }
}

// NB: Configure in a new closure so that all relocated projects have a 'shadedJar' or 'trimShadedJar' task.
configure(relocatedProjects) {
    // Make sure the main classes are ready as well when `copyShadedTestClasses` is run,
    // so that the tasks that require `copyShadedTestClasses` have everything required to run shaded tests.
    tasks.copyShadedTestClasses.configure {
        relocatedProjects.each {
            if (it.tasks.findByName('trimShadedJar')) {
                dependsOn it.tasks.trimShadedJar.path
            } else {
                dependsOn it.tasks.shadedJar.path
            }
        }
    }

    // Add tests for the shaded JAR.
    testing.suites {
        shadedTest(JvmTestSuite) {

            targets.configureEach {
                testTask.configure {

                    group = 'Verification'
                    description = 'Runs the unit tests with the shaded classes.'

                    project.ext.configureFlakyTests(it)
                    project.ext.configureCommonTestSettings(it)

                    dependsOn tasks.copyShadedTestClasses

                    // The tests against the shaded artifacts should run after the tests against the unshaded ones.
                    shouldRunAfter tasks.test

                    testClassesDirs = files(tasks.copyShadedTestClasses.destinationDir)
                    classpath = testClassesDirs

                    project.ext.relocations.each {
                        exclude "${it['to'].replace('.', '/')}/**"
                    }
                }
            }
        }
    }
    tasks.check.dependsOn testing.suites.shadedTest

    // Configurations shouldn't be created in afterEvaluate so create it eagerly here before configuring it.

    // In Gradle 8, dependencies can't be declared in a resolvable (canBeResolved=true) configuration. As a
    // workaround two configurations are created to separate the declaration and resolving roles.
    // - Directly declares dependencies to `shadedJarTestImplementation`.
    // - Indirectly resolves the dependencies of `shadedJarTestImplementation` by creating a new configuration
    //   and extends `shadedJarTestImplementation`.
    // https://discuss.gradle.org/t/problem-with-compileclasspath-after-gradle-8-update/44940
    project.configurations.create("shadedJarTestImplementation") {
        canBeResolved = false
        canBeConsumed = false
    }
    project.configurations.create("shadedJarTestRuntime") {
        canBeResolved = true
        canBeConsumed = false
        extendsFrom(project.configurations.shadedJarTestImplementation)
    }

    // Update the classpath of the 'shadedTest' task after all shaded projects are evaluated
    // so that we get the complete dependency list.
    project.afterEvaluate {
        if (numConfiguredRelocatedProjects.incrementAndGet() == relocatedProjects.size()) {
            relocatedProjects.each { p ->
                configureShadedTestImplementConfiguration(p)
            }
            relocatedProjects.each { p ->
                p.testing.suites.shadedTest.targets.configureEach {
                    testTask.configure {
                        classpath += p.files(p.configurations.getByName('shadedJarTestRuntime').resolve())
                    }
                }
            }
        }
    }

    gradle.taskGraph.whenReady {
        // Skip unshaded tests if shaded tests will run.
        if (gradle.taskGraph.hasTask(tasks.shadedTest)) {
            tasks.test.onlyIf { false }
        }
    }
}

private void configureShadowTask(Project project, ShadowJar task, boolean isMain) {
    List<Map<String, String>> relocations = project.ext.relocations

    task.configure {
        from(project.sourceSets[isMain ? 'main' : 'test'].output) {
            exclude 'META-INF/maven/**'
            // Prevent the shadow plugin from exploding the JARs in the resources directory.
            // e.g. WEB-INF/lib/hello.jar
            exclude '**/*.jar'
        }

        configurations = [project.configurations[isMain ? 'compileClasspath' : 'testCompileClasspath']]

        relocations.each { props ->
            task.relocate props['from'], props['to']
        }

        dependencies {
            if (project.hasFlags('shade')) {
                // Shade the relocated dependencies only.
                exclude(dependency({ dep ->
                    if (!relocations.find { dep.name.startsWith("${it['name']}:") }) {
                        // Do not shade the dependencies not listed in 'relocations'.
                        return true
                    }

                    if (isMain) {
                        return false
                    } else {
                        // Do not shade the dependencies which is already shaded for 'main'.
                        return project.configurations.compileClasspath.allDependencies.find { compileDep ->
                            compileDep.group == dep.moduleGroup && compileDep.name == dep.moduleName
                        }
                    }
                }))
            } else { // hasFlags('relocate')
                // Only want to rewrite source references, not bundle dependencies.
                exclude(dependency({ true }))
            }
        }
    }
}


/**
 * Finds the dependencies of {@code project} recursively and adds the found dependencies to
 * the configuration named as {@code 'shadedJarTestImplementation'}.
 */
private Configuration configureShadedTestImplementConfiguration(
        Project project, Project recursedProject = project,
        Set<ExcludeRule> excludeRules = new HashSet<>(),
        Set<Project> visitedProjects = new HashSet<>()) {

    def shadedJarTestImplementation = project.configurations.getByName('shadedJarTestImplementation')

    if (visitedProjects.contains(recursedProject)) {
        return shadedJarTestImplementation
    } else {
        visitedProjects.add(recursedProject);
    }

    if (recursedProject.tasks.findByName('trimShadedJar')) {
        project.dependencies.add(shadedJarTestImplementation.name, files(recursedProject.tasks.trimShadedJar.outJarFiles))
    } else if (recursedProject.tasks.findByName('shadedJar')) {
        project.dependencies.add(shadedJarTestImplementation.name, files(recursedProject.tasks.shadedJar.archivePath))
    }

    def shadedDependencyNames = project.ext.relocations.collect { it['name'] }
    def targetConfigurationName
    if (project == recursedProject) {
        targetConfigurationName = 'testRuntimeClasspath'
    }  else {
        targetConfigurationName = 'runtimeClasspath'
    }
    def projectDependencies = []
    recursedProject.configurations.findAll { cfg ->
        cfg.name == targetConfigurationName
    }.each { cfg ->
        cfg.allDependencies.each { dep ->
            if (dep instanceof ProjectDependency) {
                // Project dependency - recurse later.
                // Note that we recurse later to have immediate module dependencies higher precedence.
                projectDependencies.add(dep)
            } else {
                // Module dependency - add.
                if (shadedDependencyNames.contains("${dep.group}:${dep.name}")) {
                    // Skip the shaded dependencies.
                    return
                }

                if (excludeRules.find { rule ->
                    return rule.group == dep.group && rule.module == dep.name
                }) {
                    // Skip the excluded dependencies.
                    return
                }
                // Do not use `project.dependencies.add(name, dep)` that discards the classifier of
                // a dependency. See https://github.com/gradle/gradle/issues/23096
                project.configurations.getByName(shadedJarTestImplementation.name).dependencies.add(dep)
            }
        }
    }

    // Recurse into the project dependencies.
    projectDependencies.each { ProjectDependency dep ->
        configureShadedTestImplementConfiguration(
                project, dep.dependencyProject,
                excludeRules + dep.excludeRules, visitedProjects)
    }

    return shadedJarTestImplementation
}
