/*
 * Copyright 2010-2017 JetBrains s.r.o.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.jetbrains.kotlin.kapt.idea

import com.android.tools.idea.gradle.project.model.AndroidModuleModel
import com.intellij.openapi.diagnostic.Logger
import com.intellij.openapi.externalSystem.model.DataNode
import com.intellij.openapi.externalSystem.model.ProjectKeys
import com.intellij.openapi.externalSystem.model.project.*
import org.gradle.api.Plugin
import org.gradle.api.Project
import org.gradle.tooling.model.idea.IdeaModule
import org.jetbrains.kotlin.gradle.AbstractKotlinGradleModelBuilder
import org.jetbrains.kotlin.idea.framework.GRADLE_SYSTEM_ID
import org.jetbrains.plugins.gradle.model.data.GradleSourceSetData
import org.jetbrains.plugins.gradle.service.project.AbstractProjectResolverExtension
import org.jetbrains.plugins.gradle.tooling.ErrorMessageBuilder
import java.io.File
import java.io.Serializable
import java.lang.Exception
import java.lang.reflect.Modifier

interface KaptSourceSetModel : Serializable {
    val sourceSetName: String
    val isTest: Boolean
    val generatedSourcesDir: String
    val generatedClassesDir: String
    val generatedKotlinSourcesDir: String

    val generatedSourcesDirFile get() = generatedSourcesDir.takeIf { it.isNotEmpty() }?.let(::File)
    val generatedClassesDirFile get() = generatedClassesDir.takeIf { it.isNotEmpty() }?.let(::File)
    val generatedKotlinSourcesDirFile get() = generatedKotlinSourcesDir.takeIf { it.isNotEmpty() }?.let(::File)
}

class KaptSourceSetModelImpl(
        override val sourceSetName: String,
        override val isTest: Boolean,
        override val generatedSourcesDir: String,
        override val generatedClassesDir: String,
        override val generatedKotlinSourcesDir: String
) : KaptSourceSetModel

interface KaptGradleModel : Serializable {
    val isEnabled: Boolean
    val buildDirectory: File
    val sourceSets: List<KaptSourceSetModel>
}

class KaptGradleModelImpl(
        override val isEnabled: Boolean,
        override val buildDirectory: File,
        override val sourceSets: List<KaptSourceSetModel>
) : KaptGradleModel

@Suppress("unused")
class KaptProjectResolverExtension : AbstractProjectResolverExtension() {
    private companion object {
        private val LOG = Logger.getInstance(KaptProjectResolverExtension::class.java)
    }

    override fun getExtraProjectModelClasses() = setOf(KaptGradleModel::class.java)
    override fun getToolingExtensionsClasses() = setOf(KaptModelBuilderService::class.java, Unit::class.java)

    override fun populateModuleExtraModels(gradleModule: IdeaModule, ideModule: DataNode<ModuleData>) {
        val kaptModel = resolverCtx.getExtraProject(gradleModule, KaptGradleModel::class.java) ?: return

        if (kaptModel.isEnabled) {
            for (sourceSet in kaptModel.sourceSets) {
                populateAndroidModuleModelIfNeeded(ideModule, sourceSet)

                val sourceSetDataNode = ideModule.findGradleSourceSet(sourceSet.sourceSetName) ?: continue

                fun addSourceSet(path: String, type: ExternalSystemSourceType) {
                    val contentRootData = ContentRootData(GRADLE_SYSTEM_ID, path)
                    contentRootData.storePath(type, path)
                    sourceSetDataNode.createChild(ProjectKeys.CONTENT_ROOT, contentRootData)
                }

                val sourceType = if (sourceSet.isTest) ExternalSystemSourceType.TEST_GENERATED else ExternalSystemSourceType.SOURCE_GENERATED
                sourceSet.generatedSourcesDirFile?.let { addSourceSet(it.absolutePath, sourceType) }
                sourceSet.generatedKotlinSourcesDirFile?.let { addSourceSet(it.absolutePath, sourceType) }

                sourceSet.generatedClassesDirFile?.let { generatedClassesDir ->
                    val libraryData = LibraryData(GRADLE_SYSTEM_ID, "kaptGeneratedClasses")
                    libraryData.addPath(LibraryPathType.BINARY, generatedClassesDir.absolutePath)
                    val libraryDependencyData = LibraryDependencyData(sourceSetDataNode.data, libraryData, LibraryLevel.MODULE)
                    sourceSetDataNode.createChild(ProjectKeys.LIBRARY_DEPENDENCY, libraryDependencyData)
                }
            }
        }

        super.populateModuleExtraModels(gradleModule, ideModule)
    }

    private fun populateAndroidModuleModelIfNeeded(ideModule: DataNode<ModuleData>, sourceSet: KaptSourceSetModel) {
        ideModule.findAndroidModuleModel()?.let { androidModelAny ->
            // We can cast to AndroidModuleModel cause we already checked in findAndroidModuleModel() that the class exists

            val generatedKotlinSources = sourceSet.generatedKotlinSourcesDirFile ?: return

            val androidModel = androidModelAny.data as? AndroidModuleModel ?: return
            val variant = androidModel.findVariantByName(sourceSet.sourceSetName) ?: return

            androidModel.registerExtraGeneratedSourceFolder(generatedKotlinSources)

            // TODO remove this when IDEA eventually migrate to the newer Android plugin
            try {
                variant.mainArtifact.generatedSourceFolders += generatedKotlinSources
            } catch (e: Throwable) {
                // There was an error being thrown here, but the code above doesn't work for the newer versions of Android Studio 3
                // (generatedSourceFolders returns a wrapped unmodifiable list), and the thrown exception breaks the import.
                // The error will be moved back when I find a work-around for AS3.
            }
        }
    }

    private fun DataNode<ModuleData>.findAndroidModuleModel(): DataNode<*>? {
        val modelClassName = "com.android.tools.idea.gradle.project.model.AndroidModuleModel"
        val node = children.firstOrNull { it.key.dataType == modelClassName } ?: return null
        return if (!hasClassInClasspath(modelClassName)) null else node
    }

    private fun hasClassInClasspath(name: String): Boolean {
        return try {
            Class.forName(name) != null
        } catch (thr: Throwable) {
            false
        }
    }

    private fun DataNode<ModuleData>.findGradleSourceSet(sourceSetName: String): DataNode<GradleSourceSetData>? {
        val moduleName = data.id
        for (child in children) {
            val gradleSourceSetData = child.data as? GradleSourceSetData ?: continue
            if (gradleSourceSetData.id == "$moduleName:$sourceSetName") {
                @Suppress("UNCHECKED_CAST")
                return child as DataNode<GradleSourceSetData>?
            }
        }

        return null
    }
}

class KaptModelBuilderService : AbstractKotlinGradleModelBuilder() {
    override fun getErrorMessageBuilder(project: Project, e: Exception): ErrorMessageBuilder {
        return ErrorMessageBuilder.create(project, e, "Gradle import errors")
                .withDescription("Unable to build kotlin-kapt plugin configuration")
    }

    override fun canBuild(modelName: String?): Boolean = modelName == KaptGradleModel::class.java.name

    override fun buildAll(modelName: String?, project: Project): Any {
        val kaptPlugin: Plugin<*>? = project.plugins.findPlugin("kotlin-kapt")
        val kaptIsEnabled = kaptPlugin != null

        val sourceSets = mutableListOf<KaptSourceSetModel>()

        if (kaptIsEnabled) {
            project.getAllTasks(false)[project]?.forEach { compileTask ->
                if (compileTask.javaClass.name !in kotlinCompileTaskClasses) return@forEach

                val sourceSetName = compileTask.getSourceSetName()
                val isTest = sourceSetName.toLowerCase().endsWith("test")

                val kaptGeneratedSourcesDir = getKaptDirectory("getKaptGeneratedSourcesDir", project, sourceSetName)
                val kaptGeneratedClassesDir = getKaptDirectory("getKaptGeneratedClassesDir", project, sourceSetName)
                val kaptGeneratedKotlinSourcesDir = getKaptDirectory("getKaptGeneratedKotlinSourcesDir", project, sourceSetName)
                sourceSets += KaptSourceSetModelImpl(
                        sourceSetName, isTest, kaptGeneratedSourcesDir, kaptGeneratedClassesDir, kaptGeneratedKotlinSourcesDir)
            }
        }

        return KaptGradleModelImpl(kaptIsEnabled, project.buildDir, sourceSets)
    }

    private fun getKaptDirectory(funName: String, project: Project, sourceSetName: String): String {
        val kotlinKaptPlugin = project.plugins.findPlugin("kotlin-kapt") ?: return ""

        val targetMethod = kotlinKaptPlugin::class.java.methods.firstOrNull {
            Modifier.isStatic(it.modifiers) && it.name == funName && it.parameterCount == 2
        } ?: return ""

        return (targetMethod(null, project, sourceSetName) as? File)?.absolutePath ?: ""
    }
}