package eu.dnetlib.dhp.oa.graph.resolution

import com.fasterxml.jackson.databind.ObjectMapper
import eu.dnetlib.dhp.application.ArgumentApplicationParser
import eu.dnetlib.dhp.common.HdfsSupport
import eu.dnetlib.dhp.schema.common.EntityType
import eu.dnetlib.dhp.schema.oaf.{OtherResearchProduct, Publication, Result, Software, Dataset => OafDataset}
import org.apache.commons.io.IOUtils
import org.apache.hadoop.fs.{FileSystem, Path}
import org.apache.spark.SparkConf
import org.apache.spark.sql._
import org.slf4j.{Logger, LoggerFactory}

object SparkResolveEntities {

  val mapper = new ObjectMapper()
  val entities = List(EntityType.dataset, EntityType.publication, EntityType.software, EntityType.otherresearchproduct)

  def main(args: Array[String]): Unit = {
    val log: Logger = LoggerFactory.getLogger(getClass)
    val conf: SparkConf = new SparkConf()
    val parser = new ArgumentApplicationParser(IOUtils.toString(getClass.getResourceAsStream("/eu/dnetlib/dhp/oa/graph/resolution/resolve_entities_params.json")))
    parser.parseArgument(args)
    val spark: SparkSession =
      SparkSession
        .builder()
        .config(conf)
        .appName(getClass.getSimpleName)
        .master(parser.get("master")).getOrCreate()


    val graphBasePath = parser.get("graphBasePath")
    log.info(s"graphBasePath  -> $graphBasePath")
    val workingPath = parser.get("workingPath")
    log.info(s"workingPath  -> $workingPath")
    val unresolvedPath = parser.get("unresolvedPath")
    log.info(s"unresolvedPath  -> $unresolvedPath")

    val targetPath = parser.get("targetPath")
    log.info(s"targetPath  -> $targetPath")


    val fs = FileSystem.get(spark.sparkContext.hadoopConfiguration)
    fs.mkdirs(new Path(workingPath))

    resolveEntities(spark, workingPath, unresolvedPath)
    generateResolvedEntities(spark, workingPath, graphBasePath, targetPath)
  }


  def resolveEntities(spark: SparkSession, workingPath: String, unresolvedPath: String) = {
    implicit val resEncoder: Encoder[Result] = Encoders.kryo(classOf[Result])
    import spark.implicits._

    val rPid: Dataset[(String, String)] = spark.read.load(s"$workingPath/relationResolvedPid").as[(String, String)]
    val up: Dataset[(String, Result)] = spark.read.text(unresolvedPath).as[String].map(s => mapper.readValue(s, classOf[Result])).map(r => (r.getId, r))(Encoders.tuple(Encoders.STRING, resEncoder))

    rPid.joinWith(up, rPid("_2").equalTo(up("_1")), "inner").map {
      r =>
        val result = r._2._2
        val dnetId = r._1._1
        result.setId(dnetId)
        result
    }.write.mode(SaveMode.Overwrite).save(s"$workingPath/resolvedEntities")
  }


  def deserializeObject(input: String, entity: EntityType): Result = {

    entity match {
      case EntityType.publication => mapper.readValue(input, classOf[Publication])
      case EntityType.dataset => mapper.readValue(input, classOf[OafDataset])
      case EntityType.software => mapper.readValue(input, classOf[Software])
      case EntityType.otherresearchproduct => mapper.readValue(input, classOf[OtherResearchProduct])
    }
  }

  def generateResolvedEntities(spark: SparkSession, workingPath: String, graphBasePath: String, targetPath:String) = {

    implicit val resEncoder: Encoder[Result] = Encoders.kryo(classOf[Result])
    import spark.implicits._

    val re: Dataset[(String, Result)] = spark.read.load(s"$workingPath/resolvedEntities").as[Result].map(r => (r.getId, r))(Encoders.tuple(Encoders.STRING, resEncoder))
    entities.foreach {
      e => {

        val currentEntityDataset: Dataset[(String, Result)] = spark.read.text(s"$graphBasePath/$e").as[String].map(s => deserializeObject(s, e)).map(r => (r.getId, r))(Encoders.tuple(Encoders.STRING, resEncoder))

        currentEntityDataset.joinWith(re, currentEntityDataset("_1").equalTo(re("_1")), "left").map(k => {

          val a = k._1
          val b = k._2
          if (b == null)
            a._2
          else {
            a._2.mergeFrom(b._2)
            a._2
          }
        }).map(r => mapper.writeValueAsString(r))(Encoders.STRING)
          .write.mode(SaveMode.Overwrite).option("compression", "gzip").text(s"$targetPath/$e")
      }


    }
  }
}