distribution.scala

object Distribution {
  case class DistributionConfig(
    outputDirectory: File,
    configResources: Seq[File], // from src/main/resources
    distJvmOptions: String,
    distMainClass: String,
    libFilter: File ⇒ Boolean,
    additionalLibs: Seq[File])

  object DistributionKeys {
    val Dist = new Configuration("dist") describedAs "Configuration for generation module distribution"

    val outputDirectory  = settingKey[File]("Output directory for generated distribution")
    val distJvmOptions   = settingKey[String]("JVM parameters to use in start script")
    val libFilter        = settingKey[File ⇒ Boolean]("Filter of dependency jar files")
    val additionalLibs   = settingKey[Seq[File]]("Additional dependency jar files")

    val make           = taskKey[File]("Builds an application distribution folder directory")
    val pack           = taskKey[File]("Creates a distributable zip archive")

    val distMainClass  = taskKey[String]("Main class to use in start script")
    val distClean      = taskKey[Unit]("Removes Akka microkernel directory")
    val distConfig     = taskKey[DistributionConfig]("Distribution configuration")
  }

  object DistributionModule {
    import DistributionKeys._

    val distNeedsPackageBin = make <<= make.dependsOn(packageBin in Compile)

    lazy val distributionSettings: Seq[Setting[_]] = inConfig(Dist) {
      Seq(
        make                := packageBin.value,
        packageBin          := makeTask.value,
        distClean           := distCleanTask.value,
        dependencyClasspath := (dependencyClasspath in Runtime).value,
        outputDirectory     := target(_ / "app").value,
        distJvmOptions      := "-Xms1024M -Xmx1024M -Xss1M -XX:MaxPermSize=256M -XX:+UseParallelGC",
        distMainClass       := (mainClass in Compile).value.getOrElse(""),
        libFilter           := { f ⇒ true },
        additionalLibs      := defaultAdditionalLibs.value,
        distConfig          := DistributionConfig(
          outputDirectory = outputDirectory.value,
          configResources = (unmanagedResourceDirectories in Runtime).value,
          distJvmOptions  = distJvmOptions.value,
          distMainClass   = distMainClass.value,
          libFilter       = libFilter.value,
          additionalLibs  = additionalLibs.value)
      )
    } ++ Seq(make := (make in Dist).value, distNeedsPackageBin, pack := defaultPackTask.value)

    private def defaultPackTask = Def.task {
      val _      =  make.value
      val mname  = name.value
      val outDir = (outputDirectory in Dist).value
      val stream = streams.value
      stream.log.info(s"Zipping to $mname.zip")
      val mapping = (outDir ***) pair rebase(outDir, mname)
      val zip = outDir.getParentFile / s"$mname.zip"
      IO.zip(mapping, zip)
      zip
    }

    private def distCleanTask = Def.task {
      val outDir = outputDirectory.value
      val deps   = allDependencies.value
      val log    = streams.value.log
      log.info("Cleaning " + outDir)
      IO.delete(outDir)
    }

    private def makeTask = Def.task {
      import SubProjectInfo._

      val project          = thisProject.value
      val conf             = distConfig.value
      val tgt              = crossTarget.value
      val cp               = dependencyClasspath.value
      val buildStruct      = buildStructure.value
      val st               = state.value
      val log              = st.log

      val distBinPath      = conf.outputDirectory / "bin"
      val distConfigPath   = conf.outputDirectory / "config"
      val distWebApp       = conf.outputDirectory / "webapp"
      val distDeployPath   = conf.outputDirectory / "deploy"
      val distLibPath      = conf.outputDirectory / "lib"

      log.info(s"Creating distribution ${conf.outputDirectory} ...")
      IO.createDirectory(conf.outputDirectory)

      log.info("Creating launch file")
      val launchFile = new LaunchFile(conf.distJvmOptions, conf.distMainClass)
      launchFile putIntoFolder distBinPath

      log.info("Coping configuration files")
      copyConfigurationResources(conf.configResources, distConfigPath)

      log.info("Coping web components")
      copyWebComponents(conf.configResources, distWebApp)

      log.info("Coping main jar")
      copyJars(tgt, distDeployPath)

      log.info("Coping libraries")
      copyFiles(libFiles(cp, conf.libFilter), distLibPath)

      log.info("Coping additional libraries")
      copyFiles(conf.additionalLibs, distLibPath)

      log.info("Adding data from subprojects")
      val subProjectDependencies: Set[SubProjectInfo] = allSubProjectDependencies(project, buildStruct, st)
      for (subProjectDependency <- subProjectDependencies) {
        val subTarget = subProjectDependency.target
        EvaluateTask(buildStruct, packageBin in Compile, st, subProjectDependency.projectRef)
        copyJars(subTarget, distLibPath)
      }

      log.info("Distribution created.")
      conf.outputDirectory
    }

    def copyConfigurationResources(fromDirs: Seq[File], to: File) = {
      IO.createDirectory(to)
      fromDirs filter { file =>
        List("webapp").forall(!file.getName.endsWith(_))
      } foreach { IO.copyDirectory(_, to) }
    }

    def copyWebComponents(fromDirs: Seq[File], to: File) = {
      IO.createDirectory(to)
      fromDirs filterNot { file =>
        List("webapp").forall(!file.getName.endsWith(_))
      } foreach { IO.copyDirectory(_, to) }
    }

    def defaultAdditionalLibs = libraryDependencies { (libs) ⇒ Seq.empty[File] }

    private def copyJars(fromDir: File, toDir: File) = {
      val jarFiles = fromDir.listFiles.filter(f ⇒
        f.isFile &&
          f.name.endsWith(".jar") &&
          !f.name.contains("-sources") &&
          !f.name.contains("-docs"))

      copyFiles(jarFiles, toDir)
    }

    private def copyFiles(files: Seq[File], toDir: File) = {
      for (f ← files) {
        IO.copyFile(f, new File(toDir, f.getName))
      }
    }

    private def libFiles(classpath: Classpath, libFilter: File ⇒ Boolean): Seq[File] = {
      val (libs, _) = classpath.map(_.data).partition(ClasspathUtilities.isArchive)
      libs.map(_.asFile).filter(libFilter)
    }
  }

  class SubProjectInfo(
    val projectRef: ProjectRef,
    val target: File,
    val subProjects: Seq[SubProjectInfo]) {

    def recursiveSubProjects: Set[SubProjectInfo] = {
      val flatSubProjects = for {
        x ← subProjects
        y ← x.recursiveSubProjects
      } yield y
      flatSubProjects.toSet + this
    }

    override def toString: String = projectRef.project

  }

  object SubProjectInfo {
    def allSubProjectDependencies(project: ResolvedProject, buildStructure: BuildStructure, state: State): Set[SubProjectInfo] = {
      val buildUnit = buildStructure units buildStructure.root
      val uri = buildStructure.root
      val allProjects = buildUnit.defined map {
        case (id, proj) ⇒ ProjectRef(uri, id) -> proj
      }

      val subProjects: Seq[SubProjectInfo] = allProjects.collect {
        case (projRef, proj) if includeProject(proj, project) ⇒ projectInfo(projRef, proj, buildStructure, state, allProjects)
      }.toList

      val allSubProjects = subProjects.map(_.recursiveSubProjects).flatten.toSet
      allSubProjects
    }

    private def includeProject(project: ResolvedProject, parent: ResolvedProject): Boolean = {
      parent.uses.exists {
        case ProjectRef(uri, id) ⇒ id == project.id
        case _                   ⇒ false
      }
    }

    private def projectInfo(projectRef: ProjectRef, project: ResolvedProject, buildStruct: BuildStructure, state: State,
      allProjects: Map[ProjectRef, ResolvedProject]): SubProjectInfo = {

      def optionalSetting[A](key: SettingKey[A]) = key in projectRef get buildStruct.data

      def setting[A](key: SettingKey[A], errorMessage: ⇒ String) = {
        optionalSetting(key) getOrElse {
          state.log.error(errorMessage)
          throw new IllegalArgumentException()
        }
      }

      val subProjects = allProjects.collect {
        case (projRef, proj) if includeProject(proj, project) ⇒
          projectInfo(projRef, proj, buildStruct, state, allProjects)
      }.toList

      val target = setting(Keys.crossTarget, "Missing crossTarget directory")
      new SubProjectInfo(projectRef, target, subProjects)
    }
  }

  class LaunchFile(jvmOptions: String, mainClass: String) {
    private case class DistScript(name: String, content: String, executable: Boolean)

    private val distShScript = {
      s"""#!/bin/sh
       |
       |PROJECT_HOME="$$(cd "$$(cd "$$(dirname "$$0")"; pwd -P)"/..; pwd)"
       |PROJECT_CLASSPATH="$$PROJECT_HOME/config:$$PROJECT_HOME/deploy/*:$$PROJECT_HOME/webapp/*:$$PROJECT_HOME/lib/*"
       |JAVA_OPTS="$jvmOptions"
       |java $$JAVA_OPTS -cp "$$PROJECT_CLASSPATH" -Dproject.home="$$PROJECT_HOME" $mainClass
     """.toString.stripMargin
    }

    private val launchFile = DistScript("launch", distShScript, true)

    def putIntoFolder(to: File) = {
      val target = new File(to, launchFile.name)
      IO.write(target, launchFile.content)
      setExecutable(target, launchFile.executable)
    }

    private def setExecutable(target: File, executable: Boolean): Option[String] = {
      val success = target.setExecutable(executable, false)
      if (success) None else Some("Couldn't set permission on " + target)
    }
  }
}