package io.tofhir.engine.util

import io.tofhir.engine.util.FhirMappingJobFormatter.formats
import org.apache.spark.input.PortableDataStream
import org.apache.spark.sql.{Dataset, SparkSession}
import org.json4s.jackson.JsonMethods

import java.io.{BufferedReader, File, InputStreamReader, PrintWriter}
import java.net.URI
import java.nio.file.{Path, Paths}
import java.util.zip.ZipInputStream
import scala.io.Source

/**
 * Utility class that is mainly used to access and modify checkpoint-related files managed by Spark. This class serves for the purpose of being the
 * single point where Spark-related logic is implemented.
 */
object SparkUtil {

  /**
   * Gets the path of the source directory keeping the files including names of the input files.
   * This is a subdirectory managed by Spark inside the checkpoint directory.
   * The current implementation of Spark adds a root folder with name "0", keeping the other source files.
   *
   * @param sourceDirectoryRoot Root source directory created by Spark. This directory is supposed to keep job and task-specific source files.
   * @return
   */
  def getSourceDirectoryPath(sourceDirectoryRoot: Path): String = {
    Paths.get(sourceDirectoryRoot.toString, "sources", "0").toString
  }

  /**
   * Gets the path of the commit directory keeping the commit files corresponding to the sources.
   * This is a subdirectory managed by Spark inside the checkpoint directory.
   *
   * @param commitDirectory Root commit directory created by Spark. This directory is supposed to keep job and task-specific commit files.
   * @return
   */
  def getCommitDirectoryPath(commitDirectory: Path): String = {
    Paths.get(commitDirectory.toString, "commits").toString
  }

  /**
   * Write csv URI to source file
   *
   * @param sourceWriter Writer for source file
   * @param testCsvFile  Selected csv file
   */
  def writeToSourceFile(sourceWriter: PrintWriter, testCsvFile: File): Unit = {
    sourceWriter.write(s"{\"path\":\"${testCsvFile.toURI.toString.replace("\\", "\\\\")}\"}\n")
  }

  /**
   * Parses the source file corresponding to the commit number. Returns the
   *
   * @param sourceFile The source file contains the names of the input files
   * @return
   */
  def getInputFiles(sourceFile: File): Seq[File] = {
    val source = Source.fromFile(sourceFile)

    // Iterate through the lines in the file
    val sourceFiles: Seq[File] = source.getLines()
      .flatMap(line => {
        // Some lines do not contain the desired information
        try {
          // Source URI is specified in the "path" field of a json object
          val path: String = (JsonMethods.parse(line) \ "path").extract[String]
          Some(Paths.get(new URI(path)).toFile)
        } catch {
          case _: Throwable => None
        }
      })
      .toSeq
    source.close()
    sourceFiles
  }

  /**
   * Finds the latest commit file generated by the Spark. It is an integer
   *
   * @param commitFileDirectory Directory containing the commit files
   * @return
   */
  def getLastCommitOffset(commitFileDirectory: File): Int = {
    commitFileDirectory
      .listFiles()
      .filter(file => file.isFile && !file.getName.contains("."))
      .map(file => file.getName.toInt)
      .max
  }

  /**
   * Read zip file using spark context.
   *
   * @param path  Path to the zip file
   * @param spark The SparkSession
   * @return An RDD where each row represents a line from the files within the given zip file
   */
  def readZip(path: String, spark: SparkSession): Dataset[String] = {
    val fileContentRDD = spark.sparkContext.binaryFiles(path)
      .flatMap { case (name: String, content: PortableDataStream) =>
        val zis = new ZipInputStream(content.open) // We only deal with the content of the file, not with its name
        LazyList.continually(zis.getNextEntry)
          .takeWhile {
            case null => zis.close(); false
            case _ => true
          }
          .flatMap { _ => // In case there are multiple files in the ZIP, merge their lines into a single RDD
            val br = new BufferedReader(new InputStreamReader(zis))
            LazyList.continually(br.readLine()).takeWhile(_ != null) // Read the file contents line by line
          }
      }
    import spark.implicits._
    fileContentRDD.toDS() // Convert from RDD to Dataframe. Use the implicit conversions because we only have one column whose type is String
  }
}
