/*
 * 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.master.http.api.v1

import java.util
import javax.ws.rs.{BadRequestException, Consumes, GET, Path, POST, Produces}
import javax.ws.rs.core.MediaType

import scala.collection.JavaConverters._

import io.swagger.v3.oas.annotations.Operation
import io.swagger.v3.oas.annotations.media.{Content, Schema}
import io.swagger.v3.oas.annotations.responses.ApiResponse
import io.swagger.v3.oas.annotations.tags.Tag

import org.apache.celeborn.rest.v1.model.{HandleResponse, ShufflesResponse, UnregisterShufflesRequest}
import org.apache.celeborn.server.common.http.api.ApiRequestContext
import org.apache.celeborn.service.deploy.master.Master
import org.apache.celeborn.service.deploy.master.http.api.MasterHttpResourceUtils.ensureMasterIsLeader

@Tag(name = "Shuffle")
@Produces(Array(MediaType.APPLICATION_JSON))
@Consumes(Array(MediaType.APPLICATION_JSON))
class ShuffleResource extends ApiRequestContext {
  private def master = httpService.asInstanceOf[Master]
  private def statusSystem = master.statusSystem

  @Operation(description =
    "List all running shuffle keys of the service. It will return all running shuffle's key of the cluster.")
  @ApiResponse(
    responseCode = "200",
    content = Array(new Content(
      mediaType = MediaType.APPLICATION_JSON,
      schema = new Schema(implementation = classOf[ShufflesResponse]))))
  @GET
  def shuffles: ShufflesResponse = {
    val shuffles = new util.ArrayList[String]()
    statusSystem.registeredAppAndShuffles.asScala.foreach { shuffleKey =>
      val appId = shuffleKey._1
      shuffleKey._2.asScala.foreach { id =>
        shuffles.add(s"$appId-${id}")
      }
    }
    new ShufflesResponse().shuffleIds(shuffles)
  }

  @Operation(description = "Unregister shuffles from the service.")
  @ApiResponse(
    responseCode = "200",
    content = Array(new Content(
      mediaType = MediaType.APPLICATION_JSON,
      schema = new Schema(implementation = classOf[HandleResponse]))))
  @POST
  @Path("/unregister")
  def unregisterShuffles(request: UnregisterShufflesRequest): HandleResponse =
    ensureMasterIsLeader(master) {
      if (request == null) {
        throw new BadRequestException("The unregister shuffles request is required.")
      }
      val appId = normalizeParam(request.getAppId)
      val shuffleIds = request.getShuffleIds
      if (appId.isEmpty ||
        shuffleIds == null ||
        shuffleIds.isEmpty ||
        shuffleIds.asScala.exists(shuffleId => shuffleId == null || shuffleId < 0)) {
        throw new BadRequestException(
          s"appId(${request.getAppId}) is required and shuffleIds($shuffleIds) must be a nonempty list of nonnegative ids.")
      }
      val (success, message) = master.unregisterShuffles(appId, shuffleIds)
      new HandleResponse().success(success).message(message)
    }
}
