/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.celeborn.service.deploy.cluster

import java.io.File
import java.nio.file.Files
import java.util
import java.util.UUID

import scala.io.Source
import scala.util.Random

import org.junit.Assert
import org.scalatest.BeforeAndAfterAll
import org.scalatest.funsuite.AnyFunSuite

import org.apache.celeborn.client.{LifecycleManager, ShuffleClientImpl}
import org.apache.celeborn.client.read.MetricsCallback
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.identity.UserIdentifier
import org.apache.celeborn.common.internal.Logging
import org.apache.celeborn.common.protocol.CompressionCodec
import org.apache.celeborn.common.util.Utils.runCommand
import org.apache.celeborn.service.deploy.MiniClusterFeature

trait JavaCppHybridReadWriteTestBase extends AnyFunSuite
  with Logging with MiniClusterFeature with BeforeAndAfterAll {

  var masterPort = 0

  override def beforeAll(): Unit = {
    logInfo("test initialized, setup Celeborn mini cluster")
    val (m, _) = setupMiniClusterWithRandomPorts()
    masterPort = m.conf.masterPort
  }

  override def afterAll(): Unit = {
    logInfo("all test complete, stop Celeborn mini cluster")
    shutdownMiniCluster()
  }

  def testJavaWriteCppRead(codec: CompressionCodec): Unit = {
    beforeAll()
    try {
      runJavaWriteCppRead(codec)
    } finally {
      afterAll()
    }
  }

  def runJavaWriteCppRead(codec: CompressionCodec): Unit = {
    val appUniqueId = s"test-app-${UUID.randomUUID()}"
    val shuffleId = 0
    val attemptId = 0

    // Create lifecycleManager.
    val clientConf = new CelebornConf()
      .set(CelebornConf.MASTER_ENDPOINTS.key, s"localhost:$masterPort")
      .set(CelebornConf.SHUFFLE_COMPRESSION_CODEC.key, codec.name)
      .set(CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED.key, "true")
      .set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE.key, "256K")
      .set(CelebornConf.READ_LOCAL_SHUFFLE_FILE, false)
      .set("celeborn.data.io.numConnectionsPerPeer", "1")
    val lifecycleManager = new LifecycleManager(appUniqueId, clientConf)

    // Create writer shuffleClient.
    val shuffleClient =
      new ShuffleClientImpl(appUniqueId, clientConf, UserIdentifier("mock", "mock"))
    shuffleClient.setupLifecycleManagerRef(lifecycleManager.self)

    val cppResultFile = Files.createTempFile("celeborn-cpp-result-", ".txt")
    try {
      // Generate random data, write with shuffleClient and calculate result.
      val numMappers = 2
      val numPartitions = 2
      val maxData = 1000000
      val numData = 1000
      var sums = new util.ArrayList[Long](numPartitions)
      val rand = new Random()
      var prefix = "-"
      if (codec != CompressionCodec.NONE) {
        // Add duplicate strings to make the compressed length shorter than the original length.
        // Will be dropped in cpp test client.
        prefix = prefix ++ "++++++++++"
      }
      for (mapId <- 0 until numMappers) {
        for (partitionId <- 0 until numPartitions) {
          sums.add(0)
          for (i <- 0 until numData) {
            val data = rand.nextInt(maxData)
            sums.set(partitionId, sums.get(partitionId) + data)
            val dataStr = prefix + data.toString
            shuffleClient.pushOrMergeData(
              shuffleId,
              mapId,
              attemptId,
              partitionId,
              dataStr.getBytes,
              0,
              dataStr.length,
              numMappers,
              numPartitions,
              false,
              false)
          }
        }
        shuffleClient.pushMergedData(shuffleId, mapId, attemptId)
        shuffleClient.mapperEnd(shuffleId, mapId, attemptId, numMappers, numPartitions)
      }

      // Launch cpp reader to read data, calculate result and write to specific result file.
      val lifecycleManagerHost = lifecycleManager.getHost
      val lifecycleManagerPort = lifecycleManager.getPort
      val projectDirectory = new File(new File(".").getAbsolutePath)
      val cppBinRelativeDirectory = "cpp/build/celeborn/tests/"
      val cppBinFileName = "cppDataSumWithReaderClient"
      val cppBinFilePath = s"$projectDirectory/$cppBinRelativeDirectory/$cppBinFileName"
      val cppCodec = codec.name()
      // Execution command: $exec lifecycleManagerHost lifecycleManagerPort appUniqueId shuffleId attemptId numPartitions cppResultFile cppCodec
      val command = Seq(
        cppBinFilePath,
        lifecycleManagerHost,
        lifecycleManagerPort.toString,
        appUniqueId,
        shuffleId.toString,
        attemptId.toString,
        numPartitions.toString,
        cppResultFile.toString,
        cppCodec)
      println(s"run command: ${command.mkString(" ")}")
      val commandOutput = runCommand(command)
      println(s"command output: $commandOutput")

      // Verify the sum result.
      var lineCount = 0
      val source = Source.fromFile(cppResultFile.toFile, "utf-8")
      try {
        for (line <- source.getLines()) {
          val data = line.toLong
          Assert.assertEquals(data, sums.get(lineCount))
          lineCount += 1
        }
      } finally {
        source.close()
      }
      Assert.assertEquals(lineCount, numPartitions)
    } finally {
      lifecycleManager.stop()
      shuffleClient.shutdown()
      Files.deleteIfExists(cppResultFile)
    }
  }

  def testCppWriteJavaRead(codec: CompressionCodec): Unit = {
    beforeAll()
    try {
      runCppWriteJavaRead(codec)
    } finally {
      afterAll()
    }
  }

  def runCppWriteJavaRead(codec: CompressionCodec): Unit = {
    val appUniqueId = s"test-app-${UUID.randomUUID()}"
    val shuffleId = 0
    val attemptId = 0

    // Create lifecycleManager.
    val clientConf = new CelebornConf()
      .set(CelebornConf.MASTER_ENDPOINTS.key, s"localhost:$masterPort")
      .set(CelebornConf.SHUFFLE_COMPRESSION_CODEC.key, codec.name)
      .set(CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED.key, "true")
      .set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE.key, "256K")
      .set(CelebornConf.READ_LOCAL_SHUFFLE_FILE, false)
      .set("celeborn.data.io.numConnectionsPerPeer", "1")
    val lifecycleManager = new LifecycleManager(appUniqueId, clientConf)

    // Create writer shuffleClient.
    val shuffleClient =
      new ShuffleClientImpl(appUniqueId, clientConf, UserIdentifier("mock", "mock"))
    shuffleClient.setupLifecycleManagerRef(lifecycleManager.self)

    val cppResultFile = Files.createTempFile("celeborn-cpp-writer-result-", ".txt")
    try {
      val numMappers = 2
      val numPartitions = 2

      // Launch cpp writer to write data, calculate result and write to specific result file.
      val lifecycleManagerHost = lifecycleManager.getHost
      val lifecycleManagerPort = lifecycleManager.getPort
      val projectDirectory = new File(new File(".").getAbsolutePath)
      val cppBinRelativeDirectory = "cpp/build/celeborn/tests/"
      val cppBinFileName = "cppDataSumWithWriterClient"
      val cppBinFilePath = s"$projectDirectory/$cppBinRelativeDirectory/$cppBinFileName"
      val cppCodec = codec.name()
      // Execution command: $exec lifecycleManagerHost lifecycleManagerPort appUniqueId shuffleId attemptId numMappers numPartitions cppResultFile cppCodec
      val command = Seq(
        cppBinFilePath,
        lifecycleManagerHost,
        lifecycleManagerPort.toString,
        appUniqueId,
        shuffleId.toString,
        attemptId.toString,
        numMappers.toString,
        numPartitions.toString,
        cppResultFile.toString,
        cppCodec)
      println(s"run command: ${command.mkString(" ")}")
      val commandOutput = runCommand(command)
      println(s"command output: $commandOutput")

      val metricsCallback = new MetricsCallback {
        override def incBytesRead(bytesWritten: Long): Unit = {}
        override def incReadTime(time: Long): Unit = {}
      }

      var sums = new util.ArrayList[Long](numPartitions)
      for (partitionId <- 0 until numPartitions) {
        sums.add(0)
        val inputStream = shuffleClient.readPartition(
          shuffleId,
          partitionId,
          attemptId,
          0,
          0,
          Integer.MAX_VALUE,
          metricsCallback)
        try {
          var c = inputStream.read()
          var data: Long = 0
          var dataCnt = 0
          while (c != -1) {
            if (c == '-') {
              sums.set(partitionId, sums.get(partitionId) + data)
              data = 0
              dataCnt += 1
            } else {
              assert(c >= '0' && c <= '9')
              data *= 10
              data += c - '0'
            }
            c = inputStream.read()
          }
          sums.set(partitionId, sums.get(partitionId) + data)
          println(
            s"partition $partitionId sum result = ${sums.get(partitionId)}, dataCnt = $dataCnt")
        } finally {
          inputStream.close()
        }
      }

      // Verify the sum result.
      var lineCount = 0
      val source = Source.fromFile(cppResultFile.toFile, "utf-8")
      try {
        for (line <- source.getLines()) {
          val data = line.toLong
          Assert.assertEquals(data, sums.get(lineCount))
          lineCount += 1
        }
      } finally {
        source.close()
      }
      Assert.assertEquals(lineCount, numPartitions)
    } finally {
      lifecycleManager.stop()
      shuffleClient.shutdown()
      Files.deleteIfExists(cppResultFile)
    }
  }

  def testCppMergeWriteJavaRead(codec: CompressionCodec): Unit = {
    beforeAll()
    try {
      runCppMergeWriteJavaRead(codec)
    } finally {
      afterAll()
    }
  }

  def runCppMergeWriteJavaRead(codec: CompressionCodec): Unit = {
    val appUniqueId = s"test-app-${UUID.randomUUID()}"
    val shuffleId = 0
    val attemptId = 0

    val clientConf = new CelebornConf()
      .set(CelebornConf.MASTER_ENDPOINTS.key, s"localhost:$masterPort")
      .set(CelebornConf.SHUFFLE_COMPRESSION_CODEC.key, codec.name)
      .set(CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED.key, "true")
      .set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE.key, "256K")
      .set(CelebornConf.READ_LOCAL_SHUFFLE_FILE, false)
      .set("celeborn.data.io.numConnectionsPerPeer", "1")
    val lifecycleManager = new LifecycleManager(appUniqueId, clientConf)

    val shuffleClient =
      new ShuffleClientImpl(appUniqueId, clientConf, UserIdentifier("mock", "mock"))
    shuffleClient.setupLifecycleManagerRef(lifecycleManager.self)

    val cppResultFile = Files.createTempFile("celeborn-cpp-merge-writer-result-", ".txt")
    try {
      val numMappers = 2
      val numPartitions = 2

      val lifecycleManagerHost = lifecycleManager.getHost
      val lifecycleManagerPort = lifecycleManager.getPort
      val projectDirectory = new File(new File(".").getAbsolutePath)
      val cppBinRelativeDirectory = "cpp/build/celeborn/tests/"
      val cppBinFileName = "cppDataSumWithMergeWriterClient"
      val cppBinFilePath = s"$projectDirectory/$cppBinRelativeDirectory/$cppBinFileName"
      val cppCodec = codec.name()
      val command = Seq(
        cppBinFilePath,
        lifecycleManagerHost,
        lifecycleManagerPort.toString,
        appUniqueId,
        shuffleId.toString,
        attemptId.toString,
        numMappers.toString,
        numPartitions.toString,
        cppResultFile.toString,
        cppCodec)
      println(s"run command: ${command.mkString(" ")}")
      val commandOutput = runCommand(command)
      println(s"command output: $commandOutput")

      val metricsCallback = new MetricsCallback {
        override def incBytesRead(bytesWritten: Long): Unit = {}
        override def incReadTime(time: Long): Unit = {}
      }

      var sums = new util.ArrayList[Long](numPartitions)
      for (partitionId <- 0 until numPartitions) {
        sums.add(0)
        val inputStream = shuffleClient.readPartition(
          shuffleId,
          partitionId,
          attemptId,
          0,
          0,
          Integer.MAX_VALUE,
          metricsCallback)
        try {
          var c = inputStream.read()
          var data: Long = 0
          var dataCnt = 0
          while (c != -1) {
            if (c == '-') {
              sums.set(partitionId, sums.get(partitionId) + data)
              data = 0
              dataCnt += 1
            } else {
              assert(c >= '0' && c <= '9')
              data *= 10
              data += c - '0'
            }
            c = inputStream.read()
          }
          sums.set(partitionId, sums.get(partitionId) + data)
          println(
            s"partition $partitionId sum result = ${sums.get(partitionId)}, dataCnt = $dataCnt")
        } finally {
          inputStream.close()
        }
      }

      var lineCount = 0
      val source = Source.fromFile(cppResultFile.toFile, "utf-8")
      try {
        for (line <- source.getLines()) {
          val data = line.toLong
          Assert.assertEquals(data, sums.get(lineCount))
          lineCount += 1
        }
      } finally {
        source.close()
      }
      Assert.assertEquals(lineCount, numPartitions)
    } finally {
      lifecycleManager.stop()
      shuffleClient.shutdown()
      Files.deleteIfExists(cppResultFile)
    }
  }

}
