/*
 * 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.spark.sql.pipelines.graph

import java.util.concurrent.{ConcurrentHashMap, ConcurrentLinkedQueue}

import scala.jdk.CollectionConverters._

import org.apache.spark.sql.AnalysisException
import org.apache.spark.sql.catalyst.TableIdentifier
import org.apache.spark.sql.pipelines.graph.DataflowGraphTransformer.{
  TransformNodeFailedException,
  TransformNodeRetryableException
}

/**
 * Processor that is responsible for analyzing each flow and sort the nodes in
 * topological order
 */
class CoreDataflowNodeProcessor(rawGraph: DataflowGraph) {

  private val flowResolver = new FlowResolver(rawGraph)

  // Map of input identifier to resolved [[Input]].
  private val resolvedInputs = new ConcurrentHashMap[TableIdentifier, Input]()
  // Map & queue of resolved flows identifiers
  // queue is there to track the topological order while map is used to store the id -> flow
  // mapping
  private val resolvedFlowNodesMap = new ConcurrentHashMap[TableIdentifier, ResolvedFlow]()
  private val resolvedFlowNodesQueue = new ConcurrentLinkedQueue[ResolvedFlow]()

  private def processUnresolvedFlow(flow: UnresolvedFlow): ResolvedFlow = {
    val resolvedFlow = flowResolver.attemptResolveFlow(
      flow,
      rawGraph.inputIdentifiers,
      resolvedInputs.asScala.toMap
    )
    resolvedFlowNodesQueue.add(resolvedFlow)
    resolvedFlowNodesMap.put(flow.identifier, resolvedFlow)
    resolvedFlow
  }

  /**
   * Processes the node of the graph, re-arranging them if they are not topologically sorted.
   * Takes care of resolving the flows and virtualizing tables (i.e. removing tables to
   * ensure resolution is internally consistent) if needed for the nodes.
   * @param node The node to process
   * @param upstreamNodes Upstream nodes for the node
   * @return The resolved nodes generated by processing this element.
   */
  def processNode(node: GraphElement, upstreamNodes: Seq[GraphElement]): Seq[GraphElement] = {
    node match {
      case flow: UnresolvedFlow => Seq(processUnresolvedFlow(flow))
      case failedFlow: ResolutionFailedFlow => Seq(processUnresolvedFlow(failedFlow.flow))
      case table: Table =>
        // Ensure all upstreamNodes for a table are flows
        val flowsToTable = upstreamNodes.map {
          case flow: Flow => flow
          case _ =>
            throw new IllegalArgumentException(
              s"Unsupported upstream node type for table ${table.displayName}: " +
              s"${upstreamNodes.getClass}"
            )
        }
        val resolvedFlowsToTable = flowsToTable.map { flow =>
          resolvedFlowNodesMap.get(flow.identifier)
        }
        // We mark all tables as virtual to ensure resolution uses incoming flows
        // rather than previously materialized tables.
        val virtualTableInput = VirtualTableInput(
          identifier = table.identifier,
          specifiedSchema = table.specifiedSchema,
          incomingFlowIdentifiers = flowsToTable.map(_.identifier).toSet,
          availableFlows = resolvedFlowsToTable
        )
        resolvedInputs.put(table.identifier, virtualTableInput)
        Seq(table)
      case view: View =>
        // For view, add the flow to resolvedInputs and return empty.
        require(upstreamNodes.size == 1, "Found multiple flows to view")
        upstreamNodes.head match {
          case f: Flow =>
            resolvedInputs.put(view.identifier, resolvedFlowNodesMap.get(f.destinationIdentifier))
            Seq(view)
          case _ =>
            throw new IllegalArgumentException(
              s"Unsupported upstream node type for view ${view.displayName}: " +
              s"${upstreamNodes.getClass}"
            )
        }
      case sink: Sink => Seq(sink)
      case _ =>
        throw new IllegalArgumentException(s"Unsupported node type: ${node.getClass}")
    }
  }
}

private class FlowResolver(rawGraph: DataflowGraph) {

  /** Helper used to track which confs were set by which flows. */
  private case class FlowConf(key: String, value: String, flowIdentifier: TableIdentifier)

  /** Attempts resolving a single flow using the map of resolved inputs. */
  def attemptResolveFlow(
      flowToResolve: UnresolvedFlow,
      allInputs: Set[TableIdentifier],
      availableResolvedInputs: Map[TableIdentifier, Input]): ResolvedFlow = {
    val flowFunctionResult = flowToResolve.func.call(
      allInputs = allInputs,
      availableInputs = availableResolvedInputs.values.toList,
      configuration = flowToResolve.sqlConf,
      queryContext = flowToResolve.queryContext,
      queryOrigin = flowToResolve.origin
    )
    flowFunctionResult match {
        case f if f.dataFrame.isSuccess =>
          // Merge confs from any upstream views into confs for this flow.
          val allFConfs =
            (flowToResolve +:
            f.inputs.toSeq
              .map(availableResolvedInputs(_))
              .filter {
                // Input is a flow implies that the upstream table is a View.
                case _: Flow => true
                // We stop in all other cases.
                case _ => false
              }).collect {
              case g: Flow =>
                g.sqlConf.toSeq.map { case (k, v) => FlowConf(k, v, g.identifier) }
            }.flatten

          allFConfs
            .groupBy(_.key) // Key name -> Seq[FlowConf]
            .filter(_._2.length > 1) // Entries where key was set more than once
            .find(_._2.map(_.value).toSet.size > 1) // Entry where key was set with diff vals
            .foreach {
              case (key, confs) =>
                val sortedByVal = confs.sortBy(_.value)
                throw new AnalysisException(
                  "DUPLICATE_FLOW_SQL_CONF",
                  Map(
                    "key" -> key,
                    "datasetName" -> flowToResolve.displayName,
                    "flowName1" -> sortedByVal.head.flowIdentifier.unquotedString,
                    "flowName2" -> sortedByVal.last.flowIdentifier.unquotedString
                  )
                )
            }

          val newSqlConf = allFConfs.map(fc => fc.key -> fc.value).toMap
          // if the new sql confs are different from the original sql confs the flow was resolved
          // with, resolve again.
          val maybeNewFuncResult = if (newSqlConf != flowToResolve.sqlConf) {
            flowToResolve.func.call(
              allInputs = allInputs,
              availableInputs = availableResolvedInputs.values.toList,
              configuration = newSqlConf,
              queryContext = flowToResolve.queryContext,
              queryOrigin = flowToResolve.origin
            )
          } else {
            f
          }
          convertResolvedToTypedFlow(flowToResolve, maybeNewFuncResult)

        // If the flow failed due to an UnresolvedDatasetException, it means that one of the
        // flow's inputs wasn't available. After other flows are resolved, these inputs
        // may become available, so throw a retryable exception in this case.
        case f =>
          f.dataFrame.failed.toOption.collectFirst {
            case e: UnresolvedDatasetException => e
            case _ => None
          } match {
            case Some(e: UnresolvedDatasetException) =>
              throw TransformNodeRetryableException(
                e.identifier,
                new ResolutionFailedFlow(flowToResolve, flowFunctionResult)
              )
            case _ =>
              throw TransformNodeFailedException(
                new ResolutionFailedFlow(flowToResolve, flowFunctionResult)
              )
          }
      }
  }

  private def convertResolvedToTypedFlow(
      flow: UnresolvedFlow,
      funcResult: FlowFunctionResult): ResolvedFlow = {
    flow match {
      case _ if flow.once => new AppendOnceFlow(flow, funcResult)
      case _ if funcResult.dataFrame.get.isStreaming =>
        // If there's more than 1 flow to this flow's destination, we should not allow it
        // to be planned with an output mode other than Append, as the other flows will
        // then get their results overwritten.
        val mustBeAppend = rawGraph.flowsTo(flow.destinationIdentifier).size > 1
        new StreamingFlow(flow, funcResult, mustBeAppend = mustBeAppend)
      case _: UnresolvedFlow => new CompleteFlow(flow, funcResult)
    }
  }
}
