package mill
package scalalib

import java.io.ByteArrayInputStream
import java.io.ByteArrayOutputStream

import coursier.Repository
import mill.define.Task
import mill.define.TaskModule
import mill.eval.{PathRef, Result}
import mill.modules.{Assembly, Jvm}
import mill.modules.Jvm.{createAssembly, createJar}
import Lib._
import mill.scalalib.publish.{Artifact, Scope}
import mill.api.Loose.Agg

/**
  * Core configuration required to compile a single Scala compilation target
  */
trait JavaModule extends mill.Module with TaskModule with GenIdeaModule { outer =>
  def zincWorker: ZincWorkerModule = mill.scalalib.ZincWorkerModule

  trait Tests extends TestModule{
    override def moduleDeps = Seq(outer)
    override def repositories = outer.repositories
    override def javacOptions = outer.javacOptions
    override def zincWorker = outer.zincWorker
  }
  def defaultCommandName() = "run"

  def resolvePublishDependency: Task[Dep => publish.Dependency] = T.task{
    Artifact.fromDepJava(_: Dep)
  }
  def resolveCoursierDependency: Task[Dep => coursier.Dependency] = T.task{
    Lib.depToDependencyJava(_: Dep)
  }

  /**
    * Allows you to specify an explicit main class to use for the `run` command.
    * If none is specified, the classpath is searched for an appropriate main
    * class to use if one exists
    */
  def mainClass: T[Option[String]] = None

  def finalMainClassOpt: T[Either[String, String]] = T{
    mainClass() match{
      case Some(m) => Right(m)
      case None =>
        zincWorker.worker().discoverMainClasses(compile())match {
          case Seq() => Left("No main class specified or found")
          case Seq(main) => Right(main)
          case mains =>
            Left(
              s"Multiple main classes found (${mains.mkString(",")}) " +
                "please explicitly specify which one to use by overriding mainClass"
            )
        }
    }
  }

  def finalMainClass: T[String] = T{
    finalMainClassOpt() match {
      case Right(main) => Result.Success(main)
      case Left(msg)   => Result.Failure(msg)
    }
  }

  /**
    * Any ivy dependencies you want to add to this Module, in the format
    * ivy"org::name:version" for Scala dependencies or ivy"org:name:version"
    * for Java dependencies
    */
  def ivyDeps = T{ Agg.empty[Dep] }

  /**
    * Same as `ivyDeps`, but only present at compile time. Useful for e.g.
    * macro-related dependencies like `scala-reflect` that doesn't need to be
    * present at runtime
    */
  def compileIvyDeps = T{ Agg.empty[Dep] }
  /**
    * Same as `ivyDeps`, but only present at runtime. Useful for e.g.
    * selecting different versions of a dependency to use at runtime after your
    * code has already been compiled
    */
  def runIvyDeps = T{ Agg.empty[Dep] }

  /**
    * Options to pass to the java compiler
    */
  def javacOptions = T{ Seq.empty[String] }

  /** The direct dependencies of this module */
  def moduleDeps = Seq.empty[JavaModule]

  /** The direct and indirect dependencies of this module */
  def recursiveModuleDeps: Seq[JavaModule] = {
    moduleDeps.flatMap(_.transitiveModuleDeps).distinct
  }

  /** Like `recursiveModuleDeps` but also include the module itself */
  def transitiveModuleDeps: Seq[JavaModule] = {
    Seq(this) ++ recursiveModuleDeps
  }

  /**
    * Additional jars, classfiles or resources to add to the classpath directly
    * from disk rather than being downloaded from Maven Central or other package
    * repositories
    */
  def unmanagedClasspath = T{ Agg.empty[PathRef] }

  /**
    * The transitive ivy dependencies of this module and all it's upstream modules
    */
  def transitiveIvyDeps: T[Agg[Dep]] = T{
    ivyDeps() ++ Task.traverse(moduleDeps)(_.transitiveIvyDeps)().flatten
  }

  /**
    * The upstream compilation output of all this module's upstream modules
    */
  def upstreamCompileOutput = T{
    Task.traverse(recursiveModuleDeps)(_.compile)
  }

  /**
    * The transitive version of `localClasspath`
    */
  def transitiveLocalClasspath: T[Agg[PathRef]] = T{
    Task.traverse(moduleDeps)(m =>
      T.task{m.localClasspath() ++ m.transitiveLocalClasspath()}
    )().flatten
  }

  def mapDependencies = T.task{ d: coursier.Dependency => d }

  def resolveDeps(deps: Task[Agg[Dep]], sources: Boolean = false) = T.task{
    resolveDependencies(
      repositories,
      resolveCoursierDependency().apply(_),
      deps(),
      sources,
      mapDependencies = Some(mapDependencies()),
      Some(implicitly[mill.util.Ctx.Log])
    )
  }


  def repositories: Seq[Repository] = zincWorker.repositories

  /**
    * What platform suffix to use for publishing, e.g. `_sjs` for Scala.js
    * projects
    */
  def platformSuffix = T{ "" }

  private val Milestone213 = raw"""2.13.(\d+)-M(\d+)""".r

  /**
    * What shell script to use to launch the executable generated by `assembly`.
    * Defaults to a generic "universal" launcher that should work for Windows,
    * OS-X and Linux
    */
  def prependShellScript: T[String] = T{
    finalMainClassOpt().toOption match{
      case None => ""
      case Some(cls) =>
        mill.modules.Jvm.launcherUniversalScript(
          cls,
          Agg("$0"), Agg("%~dpnx0"),
          forkArgs()
        )
    }
  }

  def assemblyRules: Seq[Assembly.Rule] = Assembly.defaultRules

  /**
    * The folders where the source files for this module live
    */
  def sources = T.sources{ millSourcePath / 'src }
  /**
    * The folders where the resource files for this module live
    */
  def resources = T.sources{ millSourcePath / 'resources }
  /**
    * Folders containing source files that are generated rather than
    * hand-written; these files can be generated in this target itself,
    * or can refer to files generated from other targets
    */
  def generatedSources = T{ Seq.empty[PathRef] }

  /**
    * The folders containing all source files fed into the compiler
    */
  def allSources = T{ sources() ++ generatedSources() }

  /**
    * All individual source files fed into the compiler
    */
  def allSourceFiles = T{
    def isHiddenFile(path: os.Path) = path.last.startsWith(".")
    for {
      root <- allSources()
      if os.exists(root.path)
      path <- (if (os.isDir(root.path)) os.walk(root.path) else Seq(root.path))
      if os.isFile(path) && ((path.ext == "scala" || path.ext == "java") && !isHiddenFile(path))
    } yield PathRef(path)
  }

  /**
    * Compiles the current module to generate compiled classfiles/bytecode
    */
  def compile: T[mill.scalalib.api.CompilationResult] = T.persistent{
    zincWorker.worker().compileJava(
      upstreamCompileOutput(),
      allSourceFiles().map(_.path),
      compileClasspath().map(_.path),
      javacOptions()
    )
  }

  /**
    * The output classfiles/resources from this module, excluding upstream
    * modules and third-party dependencies
    */
  def localClasspath = T{
    resources() ++ Agg(compile().classes)
  }

  /**
    * All classfiles and resources from upstream modules and dependencies
    * necessary to compile this module
    */
  def compileClasspath = T{
    transitiveLocalClasspath() ++
    resources() ++
    unmanagedClasspath() ++
    resolveDeps(T.task{compileIvyDeps() ++ transitiveIvyDeps()})()
  }

  /**
    * All upstream classfiles and resources necessary to build and executable
    * assembly, but without this module's contribution
    */
  def upstreamAssemblyClasspath = T{
    transitiveLocalClasspath() ++
    unmanagedClasspath() ++
    resolveDeps(T.task{runIvyDeps() ++ transitiveIvyDeps()})()
  }

  /**
    * All classfiles and resources from upstream modules and dependencies
    * necessary to run this module's code after compilation
    */
  def runClasspath = T{
    localClasspath() ++
    upstreamAssemblyClasspath()
  }

  /**
    * Creates a manifest representation which can be modifed or replaced
    * The default implementation just adds the `Manifest-Version`, `Main-Class` and `Created-By` attributes
    */
  def manifest = T{
    Jvm.createManifest(finalMainClassOpt().toOption)
  }

  /**
    * Build the assembly for upstream dependencies separate from the current
    * classpath
    *
    * This should allow much faster assembly creation in the common case where
    * upstream dependencies do not change
    */
  def upstreamAssembly = T{
    createAssembly(
      upstreamAssemblyClasspath().map(_.path),
      manifest(),
      assemblyRules = assemblyRules
    )
  }

  /**
    * An executable uber-jar/assembly containing all the resources and compiled
    * classfiles from this module and all it's upstream modules and dependencies
    */
  def assembly = T{
    createAssembly(
      Agg.from(localClasspath().map(_.path)),
      manifest(),
      prependShellScript(),
      Some(upstreamAssembly().path),
      assemblyRules
    )
  }

  /**
    * A jar containing only this module's resources and compiled classfiles,
    * without those from upstream modules and dependencies
    */
  def jar = T{
    createJar(
      localClasspath().map(_.path).filter(os.exists),
      manifest()
    )
  }

  /**
   * Additional options to be used by the javadoc tool.
   * You should not set the `-d` setting for specifying the target directory,
   * as that is done in the [[docJar]] target.
   */
  def javadocOptions: T[Seq[String]] = T { Seq[String]() }

  /**
   * The documentation jar, containing all the Javadoc/Scaladoc HTML files, for
   * publishing to Maven Central
   */
  def docJar = T[PathRef] {
    val outDir = T.ctx().dest

    val javadocDir = outDir / 'javadoc
    os.makeDir.all(javadocDir)

    val files = for {
      ref <- allSources()
      if os.exists(ref.path)
      p <- (if (os.isDir(ref.path)) os.walk(ref.path) else Seq(ref.path))
      if os.isFile(p) && (p.ext == "java")
    } yield p.toNIO.toString

    val options = javadocOptions() ++ Seq("-d", javadocDir.toNIO.toString)

    if (files.nonEmpty) Jvm.runSubprocess(
      commandArgs = Seq(
        "javadoc"
      ) ++ options ++
        Seq(
          "-classpath",
          compileClasspath()
          .map(_.path)
          .filter(_.ext != "pom")
          .mkString(java.io.File.pathSeparator)
        ) ++
          files.map(_.toString),
      envArgs = Map(),
      workingDir = T.ctx().dest
    )

    createJar(Agg(javadocDir))(outDir)
  }

  /**
    * The source jar, containing only source code for publishing to Maven Central
    */
  def sourceJar = T {
    createJar((allSources() ++ resources()).map(_.path).filter(os.exists), manifest())
  }

  /**
    * Any command-line parameters you want to pass to the forked JVM under `run`,
    * `test` or `repl`
    */
  def forkArgs = T{ Seq.empty[String] }

  /**
    * Any environment variables you want to pass to the forked JVM under `run`,
    * `test` or `repl`
    */
  def forkEnv = T{ sys.env.toMap }

  /**
    * Builds a command-line "launcher" file that can be used to run this module's
    * code, without the Mill process. Useful for deployment & other places where
    * you do not want a build tool running
    */
  def launcher = T{
    Result.Success(
      Jvm.createLauncher(
        finalMainClass(),
        runClasspath().map(_.path),
        forkArgs()
      )
    )
  }

  def ivyDepsTree(inverse: Boolean = false) = T.command {
    val (flattened, resolution) = Lib.resolveDependenciesMetadata(
      repositories,
      resolveCoursierDependency().apply(_),
      transitiveIvyDeps(),
      Some(mapDependencies())
    )

    println(
      coursier.util.Print.dependencyTree(
        roots = flattened,
        resolution = resolution,
        printExclusions = false,
        reverse = inverse
      )
    )

    Result.Success()
  }

  /**
    * Runs this module's code in-process within an isolated classloader. This is
    * faster than `run`, but in exchange you have less isolation between runs
    * since the code can dirty the parent Mill process and potentially leave it
    * in a bad state.
    */
  def runLocal(args: String*) = T.command {
    Jvm.runLocal(
      finalMainClass(),
      runClasspath().map(_.path),
      args
    )
  }

  /**
    * Runs this module's code in a subprocess and waits for it to finish
    */
  def run(args: String*) = T.command{
    try Result.Success(Jvm.runSubprocess(
      finalMainClass(),
      runClasspath().map(_.path),
      forkArgs(),
      forkEnv(),
      args,
      workingDir = forkWorkingDir()
    )) catch { case e: Exception =>
       Result.Failure("subprocess failed")
    }
  }

  private[this] def backgroundSetup(dest: os.Path) = {
    val token = java.util.UUID.randomUUID().toString
    val procId = dest / ".mill-background-process-id"
    val procTombstone = dest / ".mill-background-process-tombstone"
    // The backgrounded subprocesses poll the procId file, and kill themselves
    // when the procId file is deleted. This deletion happens immediately before
    // the body of these commands run, but we cannot be sure the subprocess has
    // had time to notice.
    //
    // To make sure we wait for the previous subprocess to
    // die, we make the subprocess write a tombstone file out when it kills
    // itself due to procId being deleted, and we wait a short time on task-start
    // to see if such a tombstone appears. If a tombstone appears, we can be sure
    // the subprocess has killed itself, and can continue. If a tombstone doesn't
    // appear in a short amount of time, we assume the subprocess exited or was
    // killed via some other means, and continue anyway.
    val start = System.currentTimeMillis()
    while({
      if (os.exists(procTombstone)) {
        Thread.sleep(10)
        os.remove.all(procTombstone)
        true
      } else {
        Thread.sleep(10)
        System.currentTimeMillis() - start < 100
      }
    })()

    os.write(procId, token)
    os.write(procTombstone, token)
    (procId, procTombstone, token)
  }

  /**
    * Runs this module's code in a background process, until it dies or
    * `runBackground` is used again. This lets you continue using Mill while
    * the process is running in the background: editing files, compiling, and
    * only re-starting the background process when you're ready.
    *
    * You can also use `-w foo.runBackground` to make Mill watch for changes
    * and automatically recompile your code & restart the background process
    * when ready. This is useful when working on long-running server processes
    * that would otherwise run forever
    */
  def runBackground(args: String*) = T.command{
    val (procId, procTombstone, token) = backgroundSetup(T.ctx().dest)
    try Result.Success(Jvm.runSubprocess(
      "mill.scalalib.backgroundwrapper.BackgroundWrapper",
      (runClasspath() ++ zincWorker.backgroundWrapperClasspath()).map(_.path),
      forkArgs(),
      forkEnv(),
      Seq(procId.toString, procTombstone.toString, token, finalMainClass()) ++ args,
      workingDir = forkWorkingDir(),
      background = true
    )) catch { case e: Exception =>
       Result.Failure("subprocess failed")
    }
  }

  /**
    * Same as `runBackground`, but lets you specify a main class to run
    */
  def runMainBackground(mainClass: String, args: String*) = T.command{
    val (procId, procTombstone, token) = backgroundSetup(T.ctx().dest)
    try Result.Success(Jvm.runSubprocess(
      "mill.scalalib.backgroundwrapper.BackgroundWrapper",
      (runClasspath() ++ zincWorker.backgroundWrapperClasspath()).map(_.path),
      forkArgs(),
      forkEnv(),
      Seq(procId.toString, procTombstone.toString, token, mainClass) ++ args,
      workingDir = forkWorkingDir(),
      background = true
    )) catch { case e: Exception =>
      Result.Failure("subprocess failed")
    }
  }

  /**
    * Same as `runLocal`, but lets you specify a main class to run
    */
  def runMainLocal(mainClass: String, args: String*) = T.command {
    Jvm.runLocal(
      mainClass,
      runClasspath().map(_.path),
      args
    )
  }

  /**
    * Same as `run`, but lets you specify a main class to run
    */
  def runMain(mainClass: String, args: String*) = T.command{
    try Result.Success(Jvm.runSubprocess(
      mainClass,
      runClasspath().map(_.path),
      forkArgs(),
      forkEnv(),
      args,
      workingDir = forkWorkingDir()
    )) catch { case e: Exception =>
      Result.Failure("subprocess failed")
    }
  }

  /**
    * Override this to change the published artifact id.
    * For example, by default a scala module foo.baz might be published as foo-baz_2.12 and a java module would be foo-baz.
    * Setting this to baz would result in a scala artifact baz_2.12 or a java artifact baz.
    */
  def artifactName: T[String] = millModuleSegments.parts.mkString("-")

  /**
    * The exact id of the artifact to be published. You probably don't want to override this.
    * If you want to customize the name of the artifact, override artifactName instead.
    * If you want to customize the scala version in the artifact id, see ScalaModule.artifactScalaVersion
    */
  def artifactId: T[String] = artifactName()

  def forkWorkingDir = T{ ammonite.ops.pwd }
}

trait TestModule extends JavaModule with TaskModule {
  override def defaultCommandName() = "test"
  /**
    * What test frameworks to use.
    */
  def testFrameworks: T[Seq[String]]
  /**
    * Discovers and runs the module's tests in a subprocess, reporting the
    * results to the console
    */
  def test(args: String*) = T.command{
    val outputPath = T.ctx().dest/"out.json"

    Jvm.runSubprocess(
      mainClass = "mill.scalalib.TestRunner",
      classPath = zincWorker.scalalibClasspath().map(_.path),
      jvmArgs = forkArgs(),
      envArgs = forkEnv(),
      mainArgs =
        Seq(testFrameworks().length.toString) ++
        testFrameworks() ++
        Seq(runClasspath().length.toString) ++
        runClasspath().map(_.path.toString) ++
        Seq(args.length.toString) ++
        args ++
        Seq(outputPath.toString, T.ctx().log.colored.toString, compile().classes.path.toString, T.ctx().home.toString),
      workingDir = forkWorkingDir()
    )

    try {
      val jsonOutput = ujson.read(outputPath.toIO)
      val (doneMsg, results) = upickle.default.read[(String, Seq[TestRunner.Result])](jsonOutput)
      TestModule.handleResults(doneMsg, results)
    }catch{case e: Throwable =>
      Result.Failure("Test reporting failed: " + e)
    }

  }

  /**
    * Discovers and runs the module's tests in-process in an isolated classloader,
    * reporting the results to the console
    */
  def testLocal(args: String*) = T.command{
    val outputPath = T.ctx().dest/"out.json"

    val (doneMsg, results) = TestRunner.runTests(
      TestRunner.frameworks(testFrameworks()),
      runClasspath().map(_.path),
      Agg(compile().classes.path),
      args
    )

    TestModule.handleResults(doneMsg, results)

  }
}

object TestModule{
  def handleResults(doneMsg: String, results: Seq[TestRunner.Result]) = {

    val badTests = results.filter(x => Set("Error", "Failure").contains(x.status))
    if (badTests.isEmpty) Result.Success((doneMsg, results))
    else {
      val suffix = if (badTests.length == 1) "" else " and " + (badTests.length-1) + " more"

      Result.Failure(
        badTests.head.fullyQualifiedName + " " + badTests.head.selector + suffix,
        Some((doneMsg, results))
      )
    }
  }
}