import sbt.Keys._
import sbt._

object FixInstrumentsGeneration {

  /**
    * This task detects any changes in source files of Instruments and forces
    * recompilation of all instruments on any change. This is to ensure that the
    * Annotation Processor registers all of the instruments.
    *
    * Without that fix, incremental compilation would not register unchanged
    * instruments, leading to runtime errors.
    *
    * It should be added as a dependency of Compile / compile / compileInputs.
    */
  lazy val preCompileTask = Def.task {
    val log                 = streams.value.log
    val root                = baseDirectory.value
    val classFilesDirectory = (Compile / classDirectory).value
    val FragileFiles(fragileSources, fragileClassFiles) =
      getFragileFiles(root, classFilesDirectory)

    val fragileSourcesStore =
      streams.value.cacheStoreFactory.make("instruments_fixer")

    Tracked.diffInputs(fragileSourcesStore, FileInfo.hash)(
      fragileSources.toSet
    ) { sourcesDiff: ChangeReport[File] =>
      if (sourcesDiff.modified.nonEmpty && sourcesDiff.unmodified.nonEmpty) {
        val others =
          if (sourcesDiff.modified.size >= 2)
            s" and ${sourcesDiff.modified.size - 1} others"
          else ""
        val firstInstrument = sourcesDiff.modified.head
        val sourcesMessage  = firstInstrument.toString + others
        log.warn(
          s"Instruments sources ($sourcesMessage) have been changed.\n" +
          s"Forcing recompilation of all instruments to maintain " +
          s"consistency of generated services files."
        )

        fragileClassFiles.foreach(_.delete())
      }
    }
  }

  /**
    * This task detects if just a subset of the Instruments has been recompiled
    * (right now we did not find a way of detecting this before compilation). If
    * the Instrumentation state is detected to be inconsistent, current
    * compilation is aborted and classfiles are deleted to ensure that when
    * re-run Instrumentation will be brought back to a consistent state.
    *
    * Without that fix, incremental compilation would not register unchanged
    * instruments, leading to runtime errors.
    *
    * It should replace the default `Compile / compile` task in a project.
    */
  lazy val patchedCompile = Def.task {
    val compilationResult = (Compile / compile).value

    val log                 = streams.value.log
    val root                = baseDirectory.value
    val classFilesDirectory = (Compile / classDirectory).value
    val FragileFiles(_, fragileClassFiles) =
      getFragileFiles(root, classFilesDirectory)

    val fragileClassFilesStore =
      streams.value.cacheStoreFactory.make("instruments_classfiles")

    Tracked.diffInputs(fragileClassFilesStore, FileInfo.lastModified)(
      fragileClassFiles.toSet
    ) { sourcesDiff: ChangeReport[File] =>
      if (sourcesDiff.modified.nonEmpty && sourcesDiff.unmodified.nonEmpty) {
        fragileClassFiles.foreach(_.delete())

        val projectName = name.value
        log.error(
          "Truffle Instrumentation is not up to date, " +
          "which will lead to runtime errors\n" +
          "Fixes have been applied to ensure consistent Instrumentation state, " +
          "but compilation has to be triggered again.\n" +
          "Please re-run the previous command.\n" +
          "(If this for some reason fails, " +
          s"please do a clean build of the $projectName project)"
        )

        throw new RuntimeException("Please re-run last command")
      }
    }

    compilationResult
  }

  private case class FragileFiles(sources: Seq[File], classFiles: Seq[File])

  private def getFragileFiles(
    root: File,
    classFilesDirectory: File
  ): FragileFiles = {
    val fragileSources =
      (file(s"$root/src/main/java/") ** "*Instrument.java").get ++
      Seq(file(s"$root/src/main/java/org/enso/interpreter/Language.java"))
    val fragileClassFiles =
      (classFilesDirectory ** "*Instrument.class").get ++
      Seq(file(s"$classFilesDirectory/org/enso/interpreter/Language.class"))
    FragileFiles(fragileSources, fragileClassFiles)
  }
}