/*
 * 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.texera.amber.engine.architecture.controller.execution

import org.apache.texera.amber.core.virtualidentity.PhysicalOpIdentity
import org.apache.texera.amber.engine.architecture.controller.execution.ExecutionUtils.aggregateMetrics
import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState._
import org.apache.texera.amber.engine.architecture.scheduling.{Region, RegionIdentity}
import org.apache.texera.amber.engine.common.executionruntimestate.OperatorMetrics

import scala.collection.mutable

case class WorkflowExecution() {

  // region executions are stored with LinkedHashMap to maintain their creation order.
  private val regionExecutions: mutable.LinkedHashMap[RegionIdentity, RegionExecution] =
    mutable.LinkedHashMap()

  /**
    * Initializes or retrieves a `RegionExecution` for a given `Region`. If not already
    * initialized, it creates and returns a new `RegionExecution`; otherwise, an assertion
    * error is thrown if re-initialization is attempted.
    *
    * @param region The `Region` for which to initialize or retrieve the `RegionExecution`.
    * @return The `RegionExecution` associated with the given `Region`.
    * @throws AssertionError if the `RegionExecution` has already been initialized.
    */
  def initRegionExecution(region: Region): RegionExecution = {
    // ensure the region execution hasn't been initialized already.
    assert(
      !regionExecutions.contains(region.id),
      s"RegionExecution of ${region.id} already initialized."
    )
    regionExecutions.getOrElseUpdate(region.id, RegionExecution(region))
  }

  def restartRegionExecution(region: Region): RegionExecution = {
    regionExecutions.get(region.id).foreach { existingRegionExecution =>
      assert(
        existingRegionExecution.isCompleted,
        s"Cannot restart running RegionExecution of ${region.id}."
      )
    }
    val regionExecution = RegionExecution(region)
    regionExecutions.put(region.id, regionExecution)
    regionExecution
  }

  /**
    * Retrieves a specific `RegionExecution` by its identifier.
    *
    * @param regionId The unique identifier of the region for which the execution is to be retrieved.
    * @return The `RegionExecution` associated with the specified `regionId`.
    */
  def getRegionExecution(regionId: RegionIdentity): RegionExecution = regionExecutions(regionId)

  def hasRegionExecution(regionId: RegionIdentity): Boolean = regionExecutions.contains(regionId)

  /**
    * Retrieves all `RegionExecutions` that are currently in running state,
    * preserving the order in which they were created.
    *
    * This method filters the executions to include only those that have not completed.
    *
    * @return An `Iterable` of `RegionExecution` objects that are in running state.
    */
  def getRunningRegionExecutions: Iterable[RegionExecution] = {
    regionExecutions.values.filterNot(_.isCompleted)
  }

  /**
    * Retrieve the runtime stats of all `RegionExecutions`
    *
    * @return A `Map` with key being `Logical Operator ID` and the value being operator runtime statistics
    */
  def getAllRegionExecutionsStats: Map[String, OperatorMetrics] = {
    val allRegionExecutions: Iterable[RegionExecution] = getAllRegionExecutions

    val statsMap: Map[PhysicalOpIdentity, OperatorMetrics] = allRegionExecutions.flatMap {
      regionExecution =>
        regionExecution.getStats.map {
          case (physicalOpIdentity, operatorMetrics) =>
            (physicalOpIdentity, operatorMetrics)
        }
    }.toMap

    val aggregatedStats: Map[String, OperatorMetrics] =
      statsMap.groupBy(_._1.logicalOpId.id).map {
        case (logicalOpId, stats) =>
          (logicalOpId, aggregateMetrics(stats.values))
      }
    aggregatedStats
  }

  /**
    * Retrieves all `RegionExecutions`, preserving the order in which they were created.
    *
    * This method provides access to all executions, regardless of their state.
    *
    * @return An `Iterable` of all `RegionExecution` objects in the order they were added.
    */
  def getAllRegionExecutions: Iterable[RegionExecution] = regionExecutions.values

  /**
    * Retrieves the latest `OperatorExecution` associated with the specified physical operatorId.
    *
    * This method searches through all `RegionExecutions` in reverse creation order to find the most recent
    * `OperatorExecution` that matches the given physical operatorId. It assumes that each `RegionExecution`
    * may contain zero or exactly one `OperatorExecution` instance, and it returns the latest one found that
    * corresponds to the specified operatorId.
    *
    * @param physicalOpId The unique identifier of the physical operator for which the latest execution is
    *                     to be retrieved.
    * @return The latest `OperatorExecution` instance associated with the given physical operatorId.
    * @throws NoSuchElementException if no `OperatorExecution` is found for the specified operatorId.
    */
  def getLatestOperatorExecution(physicalOpId: PhysicalOpIdentity): OperatorExecution = {
    getLatestOperatorExecutionOption(physicalOpId).get
  }

  /**
    * Returns the latest `OperatorExecution` for a physical operator if it has been initialized.
    *
    * This is the safe counterpart of `getLatestOperatorExecution` for callers that may traverse
    * operators before their region is launched (e.g., full-graph stats queries while execution is still
    * progressing through schedule levels).
    */
  def getLatestOperatorExecutionOption(
      physicalOpId: PhysicalOpIdentity
  ): Option[OperatorExecution] = {
    regionExecutions.values.toSeq
      .findLast(regionExecution => regionExecution.hasOperatorExecution(physicalOpId))
      .map(_.getOperatorExecution(physicalOpId))
  }

  def isCompleted: Boolean = getState == WorkflowAggregatedState.COMPLETED

  def getState: WorkflowAggregatedState = {
    val regionStates = regionExecutions.values.map(_.getState)
    if (regionStates.isEmpty) {
      return WorkflowAggregatedState.UNINITIALIZED
    }
    if (regionStates.forall(_ == COMPLETED)) {
      return WorkflowAggregatedState.COMPLETED
    }
    val unCompletedOpStates = regionExecutions.values
      .filter(_.getState != COMPLETED)
      .flatMap(_.getAllOperatorExecutions.map(_._2.getState))
      .filter(_ != COMPLETED)
    if (unCompletedOpStates.forall(_ == UNINITIALIZED)) {
      return WorkflowAggregatedState.UNINITIALIZED
    }
    val runningOpStates = unCompletedOpStates.filter(_ != UNINITIALIZED)
    if (runningOpStates.exists(_ == RUNNING)) {
      WorkflowAggregatedState.RUNNING
    } else if (runningOpStates.forall(_ == PAUSED)) {
      WorkflowAggregatedState.PAUSED
    } else if (runningOpStates.forall(_ == READY)) {
      WorkflowAggregatedState.READY
    } else {
      WorkflowAggregatedState.UNKNOWN
    }
  }

}
