package wacc

import cats.data.{Chain, NonEmptyList}
import parsley.{Failure, Success}

import java.nio.file.{Files, Path}
import cats.syntax.all._

import cats.effect.IO
import cats.effect.ExitCode

import com.monovore.decline._
import com.monovore.decline.effect._

import org.typelevel.log4cats.slf4j.Slf4jLogger
import org.typelevel.log4cats.Logger

import assemblyIR as asm
import cats.data.ValidatedNel
import java.io.File
import cats.data.NonEmptySeq

/*
TODO:
  1) IO correctness
  2) Errors can be handled more gracefully - currently, parallelised compilation is not fail fast as far as I am aware
  3) splitting the file up and nicer refactoring
  4) logging could be removed
  5) general cleanup and comments (things like replacing home/<user> with ~ , and names of parameters and args, descriptions etc)
 */

private val SUCCESS = ExitCode.Success.code
private val ERROR = ExitCode.Error.code

given logger: Logger[IO] = Slf4jLogger.getLogger[IO]

val logOpt: Opts[Boolean] =
  Opts.flag("log", "Enable logging for additional compilation details", short = "l").orFalse

def validateFile(path: Path): ValidatedNel[String, Path] = {
  (for {
    // TODO: redundant 2nd parameter :(
    _ <- Either.cond(Files.exists(path), (), s"File '${path}' does not exist")
    _ <- Either.cond(Files.isRegularFile(path), (), s"File '${path}' must be a regular file")
    _ <- Either.cond(path.toString.endsWith(".wacc"), (), "File must have .wacc extension")
  } yield path).toValidatedNel
}

val filesOpt: Opts[NonEmptyList[Path]] =
  Opts.arguments[Path]("files").mapValidated {
    _.traverse(validateFile)
  }

val outputOpt: Opts[Option[Path]] =
  Opts
    .option[Path]("output", metavar = "path", help = "Output directory for compiled files.")
    .validate("Must have permissions to create & access the output path") { path =>
      try {
        Files.createDirectories(path)
        true
      } catch {
        case e: java.nio.file.AccessDeniedException =>
          false
      }
    }
    .validate("Output path must be a directory") { path =>
      Files.isDirectory(path)
    }
    .orNone

def frontend(
    contents: String,
    file: File
): IO[Either[NonEmptySeq[Error], microWacc.Program]] =
  parser.parse(contents) match {
    case Failure(msg) => IO.pure(Left(NonEmptySeq.one(Error.SyntaxError(file, msg))))
    case Success(fn) =>
      val partialProg = fn(file)

      for {
        (typedProg, errors) <- semantics.check(partialProg)
        res = NonEmptySeq.fromSeq(errors.iterator.toSeq).map(Left(_)).getOrElse(Right(typedProg))
      } yield res
  }

def backend(typedProg: microWacc.Program): Chain[asm.AsmLine] =
  asmGenerator.generateAsm(typedProg)

def compile(
    filePath: Path,
    outputDir: Option[Path],
    log: Boolean
): IO[Int] = {
  val logAction: String => IO[Unit] =
    if (log) logger.info(_)
    else (_ => IO.unit)

  def readSourceFile: IO[String] =
    IO.blocking(os.read(os.Path(filePath)))

  // TODO: path, file , the names are confusing (when Path is the type but we are working with files)
  def writeOutputFile(typedProg: microWacc.Program, outputPath: Path): IO[Unit] =
    val backendStart = System.nanoTime()
    val asmLines = backend(typedProg)
    val backendEnd = System.nanoTime()
    writer.writeTo(asmLines, outputPath) *>
      logAction(
        s"Backend time (${filePath.toRealPath()}): ${(backendEnd - backendStart).toFloat / 1e6} ms"
      ) *>
      IO.blocking(println(s"Success: ${outputPath.toRealPath()}"))

  def processProgram(contents: String, file: File, outDir: Path): IO[Int] =
    val frontendStart = System.nanoTime()
    for {
      frontendResult <- frontend(contents, file)
      frontendEnd = System.nanoTime()
      _ <- logAction(
        s"Frontend time (${filePath.toRealPath()}): ${(frontendEnd - frontendStart).toFloat / 1e6} ms"
      )
      res <- frontendResult match {
        case Left(errors) =>
          val code = errors.map(err => err.exitCode).toList.min
          val errorMsg = errors.map(formatError).toIterable.mkString("\n")
          for {
            _ <- logAction(s"Compilation failed for $filePath\nExit code: $code")
            _ <- IO.blocking(
              // Explicit println since we want this to always show without logger thread info e.t.c.
              println(s"Compilation failed for ${file.getCanonicalPath}:\n$errorMsg")
            )
          } yield code

        case Right(typedProg) =>
          val outputFile = outDir.resolve(filePath.getFileName.toString.stripSuffix(".wacc") + ".s")
          writeOutputFile(typedProg, outputFile).as(SUCCESS)
      }
    } yield res

  for {
    contents <- readSourceFile
    _ <- logAction(s"Compiling file: ${filePath.toAbsolutePath}")
    exitCode <- processProgram(contents, filePath.toFile, outputDir.getOrElse(filePath.getParent))
  } yield exitCode
}

def compileCommandParallel(
    files: NonEmptyList[Path],
    log: Boolean,
    outDir: Option[Path]
): IO[ExitCode] =
  files
    .parTraverse { file => compile(file.toAbsolutePath, outDir, log) }
    .map { exitCodes =>
      exitCodes.filter(_ != 0) match {
        case Nil        => ExitCode.Success
        case errorCodes => ExitCode(errorCodes.min)
      }
    }

object Main
    extends CommandIOApp(
      name = "wacc",
      header = "The ultimate WACC compiler",
      version = "1.0"
    ) {
  def main: Opts[IO[ExitCode]] =
    (filesOpt, logOpt, outputOpt).mapN { (files, log, outDir) =>
      compileCommandParallel(files, log, outDir)
    }

}